#include <openssl/evp.h>
#include <openssl/pem.h>
#include <stdio.h>
#include <stdlib.h>

void handleErrors(void) {
    fprintf(stderr, "An error occurred.\n");
    exit(1);
}

int main(int argc, char *argv[]) {
    if (argc != 4) {
        fprintf(stderr, "Usage: %s <pubkey.pem> <input file> <output file>\n", argv[0]);
        return 1;
    }

    char *pubkey_filename = argv[1];
    char *input_file = argv[2];
    char *output_file = argv[3];

    FILE *f_input, *f_output, *f_pubkey;
    unsigned char buffer[1024];
    unsigned char *ciphertext;
    size_t ciphertext_len;
    size_t bytes_read;

    EVP_PKEY *pubkey = NULL;
    EVP_PKEY_CTX *ctx = NULL;

    // Load public key
    if (!(f_pubkey = fopen(pubkey_filename, "r"))) {
        fprintf(stderr, "Unable to open public key file %s\n", pubkey_filename);
        return 1;
    }

    if (!(pubkey = PEM_read_PUBKEY(f_pubkey, NULL, NULL, NULL))) {
        fprintf(stderr, "Error loading public key\n");
        fclose(f_pubkey);
        return 1;
    }
    fclose(f_pubkey);

    // Initialise the library
    if (!OPENSSL_init_crypto(OPENSSL_INIT_LOAD_CRYPTO_STRINGS, NULL))
        handleErrors();

    // Create and initialise the context
    if (!(ctx = EVP_PKEY_CTX_new(pubkey, NULL)))
        handleErrors();

    if (EVP_PKEY_encrypt_init(ctx) <= 0)
        handleErrors();

    // Open files
    if (!(f_input = fopen(input_file, "rb"))) {
        fprintf(stderr, "Could not open %s for reading\n", input_file);
        return 1;
    }

    if (!(f_output = fopen(output_file, "wb"))) {
        fprintf(stderr, "Could not open %s for writing\n", output_file);
        return 1;
    }

    // Determine buffer size for output
    if (EVP_PKEY_encrypt(ctx, NULL, &ciphertext_len, buffer, sizeof(buffer)) <= 0)
        handleErrors();
    ciphertext = malloc(ciphertext_len);
    if (!ciphertext) {
        fprintf(stderr, "Memory allocation failed\n");
        return 1;
    }

    // Encrypt the data
    while ((bytes_read = fread(buffer, 1, sizeof(buffer), f_input)) > 0) {
        if (EVP_PKEY_encrypt(ctx, ciphertext, &ciphertext_len, buffer, bytes_read) <= 0)
            handleErrors();
        fwrite(ciphertext, 1, ciphertext_len, f_output);
    }

    // Clean up
    EVP_PKEY_free(pubkey);
    EVP_PKEY_CTX_free(ctx);
    fclose(f_input);
    fclose(f_output);
    free(ciphertext);

    printf("Encryption complete.\n");

    return 0;
}
